package demo.db.zookeeper.kafka;

import com.alibaba.fastjson.JSONObject;
import demo.db.zookeeper.ZkUtils;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.ZooKeeper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.io.IOException;
import java.util.List;

/**
 * 获取kafka在zk上的注册信息
 * 
 * @author junying.han
 *
 */
public class KafkaZookeeperClient implements Closeable {

    private static Logger logger = LoggerFactory.getLogger(KafkaZookeeperClient.class);

    /** 访问kafka zk注册中心的客户端 */
    private ZooKeeper zkClient;

    public KafkaZookeeperClient(ZooKeeper kafkaZkClient) throws IOException, InterruptedException {
        if (kafkaZkClient == null) {
            throw new IllegalArgumentException("kafkaZkClient can't be null");
        }
        this.zkClient = kafkaZkClient;
    }

    public KafkaZookeeperClient(String kafkaZkAddress) throws IOException, InterruptedException {
        if (kafkaZkAddress == null || kafkaZkAddress.isEmpty()) {
            throw new IllegalArgumentException("kafkaZkAddress can't be null");
        }
        this.zkClient = ZkUtils.buildZkClient(kafkaZkAddress, 9000);
    }

    public static final String BrokerTopicPartitionsPattern = "/brokers/topics/%s/partitions";

    public static final String BrokerTopicPartitionStatePattern = "/brokers/topics/%s/partitions/%s/state";

    /**
     * get某brokerId的信息，比如IP、端口等
     * 
     * @param brokerId
     * @return
     * @throws KeeperException
     * @throws InterruptedException
     */
    public JSONObject brokerInfo(int brokerId) throws KeeperException, InterruptedException {
        String path = BrokerIdsPath + "/" + brokerId;
        String data = getDataAsString(path);
        return JSONObject.parseObject(data);
    }

    /**
     * 返回某个topic的信息
     * 
     * @param topic
     * @return e.g. {"version":1,"partitions":{"2":[0],"1":[2],"3":[2],"0":[0]}}
     * @throws InterruptedException
     * @throws KeeperException
     */
    public JSONObject topicInfo(String topic) throws KeeperException, InterruptedException {
        String path = BrokerTopicsPath + "/" + topic;
        String data = getDataAsString(path);
        return JSONObject.parseObject(data);
    }

    /**
     * 返回某个topic的分区信息
     * 
     * @param topic
     * @return e.g. {"controller_epoch": 1,"leader": 2,"version": 1,"leader_epoch": 0,"isr": [2, 1]}
     * @throws InterruptedException
     * @throws KeeperException
     */
    public JSONObject topicPatitionState(String topic, int partition) throws KeeperException, InterruptedException {
        String path = String.format(BrokerTopicPartitionStatePattern, topic, partition);
        String data = getDataAsString(path);
        return JSONObject.parseObject(data);
    }

    /**
     * 获取某topic下的所有partitions
     * 
     * @param topic
     * @return
     * @throws KeeperException
     * @throws InterruptedException
     */
    public List<String> topicPartitions(String topic) throws KeeperException, InterruptedException {
        String path = String.format(BrokerTopicPartitionsPattern, topic);
        return zkClient.getChildren(path, false);
    }

    /**
     * 列出未删除的topic
     * 
     * @return
     * @throws KeeperException
     * @throws InterruptedException
     */
    public List<String> listTopics() throws KeeperException, InterruptedException {
        long t0 = System.currentTimeMillis();
        List<String> list = zkClient.getChildren(BrokerTopicsPath, false);
        logger.info("list topic cost {} ms.", System.currentTimeMillis() - t0);
        return list;
    }

    /**
     * 返回某个路径的数据
     * 
     * @param path
     * @return
     * @throws KeeperException
     * @throws InterruptedException
     */
    public String getDataAsString(String path) throws KeeperException, InterruptedException {
        byte[] data = zkClient.getData(path, false, null);
        if (data == null || data.length == 0) {
            return null;
        }
        return new String(data);
    }

    public static final String AdminPath = "/admin";
    public static final String ReassignPartitionsPath = "/admin/reassign_partitions";
    public static final String DeleteTopicsPath = "/admin/delete_topics";
    public static final String PreferredReplicaLeaderElectionPath = "/admin/preferred_replica_election";
    public static final String BrokersPath = "/brokers";
    public static final String BrokerIdsPath = "/brokers/ids";
    public static final String BrokerTopicsPath = "/brokers/topics";
    public static final String BrokerSequenceIdPath = "/brokers/seqid";
    public static final String ClusterPath = "/cluster";
    public static final String ClusterIdPath = "/cluster/id";
    public static final String ConfigPath = "/config";
    public static final String ConfigChangesPath = "/config/changes";
    public static final String ConfigUsersPath = "/config/users";
    public static final String ConfigBrokersPath = "/config/brokers";
    public static final String ControllerPath = "/controller";
    public static final String ControllerEpochPath = "/controller_epoch";
    public static final String IsrChangeNotificationPath = "/isr_change_notification";
    public static final String LogDirEventNotificationPath = "/log_dir_event_notification";
    public static final String KafkaAclPath = "/kafka-acl";
    public static final String KafkaAclChangesPath = "/kafka-acl-changes";
    public static final String ConsumersPath = "/consumers";
    public static final String ProducerIdBlockPath = "/latest_producer_id_block";

    @Override
    public void close() throws IOException {
        if (zkClient != null) {
            try {
                zkClient.close();
            } catch (InterruptedException e) {
                logger.error("", e);
            }
        }

    }

}
